sst5 support instead of sst2 #1
Conversation
|
Sorry - could you give more details here? 😅 What's "sst2"? What's "sst5"? |
I added some such info in the description. |
| @@ -141,8 +142,10 @@ impl SentimentModel { | |||
| let labels = self.sequence_classification_model.predict(input); | |||
| let mut sentiments = Vec::with_capacity(labels.len()); | |||
| for label in labels { | |||
| let polarity = if label.id == 1 { | |||
| let polarity = if label.id == 4 || label.id == 3 { | |||
There was a problem hiding this comment.
Okay so this basically gains us sst5 but loses us sst2 but that's fine because the current production model we're using is sst5. And this is why, as you mention in the PR description, it'd be great to get this into rust-bert as something configurable per model so the sentiment pipeline works with whatever IDs the underlying model will return?
There was a problem hiding this comment.
Yes - and I think it would be a fun rust-bert contribution to try, but maybe tricky.
There was a problem hiding this comment.
Yeah I wonder how these models work and if we could get "lucky" with a "specify the neutral range and anything above that is positive and anything below is negative". But it's technically more flexible to require the mapping as you say. Yeah interesting to try and figure out.
There was a problem hiding this comment.
Well, the models often come with something like config.json which explicitly contains the mapping.
Thanks for double checking this - one question will remain, should this actually get merged? The PR is useful to get a review, but I imagine it's easier to keep this in a branch which we point to. Is there a canonical preference? |
Yeah that is a good question. I feel like normally the process is:
but here we're like 100% sure that the thing we PR to Might be best to dismiss my approval and add commits to this branch to approach the thing we think we'd submit to I think we could also merge this and other PRs into our |
For sentiment, sst5 predicts 5 sentiment categories - "very positive", "positive", "neutral", "negative", and "very negative".
sst2 predicts 2 sentiment categories - "positive" and "negative".
rust-bert only supports sst2, not sst5, but we want to use sst5
This is a somewhat ugly, though not that bad imo hot-fix - the proper solution might be to work on getting rust-bert to support both, or rather to support reading a config file with a label -> sentiment mapping (although that would remove the ability to have the sentiment values be in an enum) - I don't see how to do that without making a breaking API change?
This also is added on top of a particular commit - I was not able to easily update our rust-bert to 0.19.